package org.wmq.taskschedule.util;

import org.opencv.core.Mat;
import org.opencv.imgcodecs.Imgcodecs;
import org.springframework.web.multipart.MultipartFile;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.awt.image.DataBufferByte;
import java.io.*;
import java.nio.ByteBuffer;
import java.util.Base64;

/**
 * 图像处理工具类
 * 提供MultipartFile与Mat类型之间的转换功能
 */
public class ImageConversionUtil {

    /**
     * 将MultipartFile转换为OpenCV的Mat对象
     * 
     * @param file 待转换的MultipartFile
     * @return 转换后的Mat对象
     * @throws IOException 转换过程中的IO异常
     */
    public static Mat convertMultipartFileToMat(MultipartFile file) throws IOException {
        if (file == null || file.isEmpty()) {
            throw new IllegalArgumentException("文件不能为空");
        }

        // 获取文件名，用于确定图像类型
        String fileName = file.getOriginalFilename();
        if (fileName == null) {
            throw new IllegalArgumentException("文件名不能为空");
        }

        // 读取文件内容到字节数组
        byte[] fileBytes = file.getBytes();

        // 创建临时Mat对象
        Mat mat = Imgcodecs.imdecode(new org.opencv.core.MatOfByte(fileBytes), Imgcodecs.IMREAD_UNCHANGED);

        if (mat.empty()) {
            throw new IOException("无法解码图像文件: " + fileName);
        }

        return mat;
    }

    /**
     * 将Mat对象转换为字节数组
     * 
     * @param mat 待转换的Mat对象
     * @param imageFormat 图像格式（如"jpg"、"png"等）
     * @return 转换后的字节数组
     * @throws IOException 转换过程中的IO异常
     */
    public static byte[] convertMatToByteArray(Mat mat, String imageFormat) throws IOException {
        if (mat == null || mat.empty()) {
            throw new IllegalArgumentException("Mat对象不能为空");
        }

        if (imageFormat == null || imageFormat.isEmpty()) {
            imageFormat = "jpg"; // 默认使用jpg格式
        }

        // 创建临时MatOfByte对象用于存储图像数据
        org.opencv.core.MatOfByte matOfByte = new org.opencv.core.MatOfByte();

        // 将Mat编码为指定格式的字节数组
        Imgcodecs.imencode("." + imageFormat.toLowerCase(), mat, matOfByte);

        // 转换为Java字节数组
        return matOfByte.toArray();
    }

    /**
     * 将Mat对象转换为BufferedImage对象
     * 
     * @param mat 待转换的Mat对象
     * @return 转换后的BufferedImage对象
     */
    public static BufferedImage convertMatToBufferedImage(Mat mat) {
        if (mat == null || mat.empty()) {
            throw new IllegalArgumentException("Mat对象不能为空");
        }

        int type = BufferedImage.TYPE_BYTE_GRAY;
        if (mat.channels() > 1) {
            type = BufferedImage.TYPE_3BYTE_BGR;
        }

        BufferedImage image = new BufferedImage(mat.cols(), mat.rows(), type);
        byte[] data = ((DataBufferByte) image.getRaster().getDataBuffer()).getData();
        mat.get(0, 0, data);

        return image;
    }

    /**
     * 将BufferedImage转换为Mat对象
     * 
     * @param image 待转换的BufferedImage对象
     * @return 转换后的Mat对象
     */
    public static Mat convertBufferedImageToMat(BufferedImage image) {
        if (image == null) {
            throw new IllegalArgumentException("BufferedImage对象不能为空");
        }

        Mat mat = new Mat(image.getHeight(), image.getWidth(), org.opencv.core.CvType.CV_8UC3);
        byte[] data = ((DataBufferByte) image.getRaster().getDataBuffer()).getData();
        mat.put(0, 0, data);

        return mat;
    }

    /**
     * 将Mat对象转换为Base64编码的字符串
     * 
     * @param mat 待转换的Mat对象
     * @param imageFormat 图像格式
     * @return Base64编码的图像字符串
     * @throws IOException 转换过程中的IO异常
     */
    public static String convertMatToBase64(Mat mat, String imageFormat) throws IOException {
        byte[] imageBytes = convertMatToByteArray(mat, imageFormat);
        return Base64.getEncoder().encodeToString(imageBytes);
    }

    /**
     * 将Base64编码的字符串转换为Mat对象
     * 
     * @param base64Image Base64编码的图像字符串
     * @return 转换后的Mat对象
     * @throws IOException 转换过程中的IO异常
     */
    public static Mat convertBase64ToMat(String base64Image) throws IOException {
        if (base64Image == null || base64Image.isEmpty()) {
            throw new IllegalArgumentException("Base64字符串不能为空");
        }

        // 移除Base64字符串中的前缀（如果有）
        String pureBase64 = base64Image;
        if (base64Image.contains(",")) {
            pureBase64 = base64Image.split(",")[1];
        }

        // 解码Base64字符串到字节数组
        byte[] imageBytes = Base64.getDecoder().decode(pureBase64);

        // 将字节数组转换为Mat对象
        return Imgcodecs.imdecode(new org.opencv.core.MatOfByte(imageBytes), Imgcodecs.IMREAD_UNCHANGED);
    }

    /**
     * 将MultipartFile保存为临时文件
     * 
     * @param file 待保存的MultipartFile
     * @return 临时文件对象
     * @throws IOException 保存过程中的IO异常
     */
    public static File convertMultipartFileToFile(MultipartFile file) throws IOException {
        if (file == null || file.isEmpty()) {
            throw new IllegalArgumentException("文件不能为空");
        }

        File tempFile = File.createTempFile("temp_image_", "." + getFileExtension(file.getOriginalFilename()));
        try (FileOutputStream fos = new FileOutputStream(tempFile)) {
            fos.write(file.getBytes());
        }
        
        // 当JVM退出时删除临时文件
        tempFile.deleteOnExit();
        
        return tempFile;
    }

    /**
     * 获取文件扩展名
     * 
     * @param fileName 文件名
     * @return 文件扩展名（小写）
     */
    private static String getFileExtension(String fileName) {
        if (fileName == null || !fileName.contains(".")) {
            return "jpg"; // 默认返回jpg
        }
        return fileName.substring(fileName.lastIndexOf(".") + 1).toLowerCase();
    }
}